package org.zjx.message;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONObject;
import com.alibaba.fastjson2.JSONWriter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.zjx.message.messages.*;

import java.lang.reflect.Type;
import java.nio.charset.StandardCharsets;
import java.util.Map;
import java.util.ServiceLoader;
import java.util.concurrent.ConcurrentHashMap;

public class MessageFactory {
    private static final Logger logger = LoggerFactory.getLogger(MessageFactory.class);
    private static final Map<String, Class<? extends IMessage>> messageLib = new ConcurrentHashMap<>();

    static {
        // 使用ServiceLoader自动发现所有指令提供者
        ServiceLoader<IMessage> loader = ServiceLoader.load(IMessage.class);
        for (IMessage provider : loader) {
            messageLib.put(provider.getClass().getName(), provider.getClass());
        }

        // 确保核心指令总是可用
        registerMessageType(AckMessage.class);
        registerMessageType(CommandMessage.class);
        registerMessageType(HeartbeatMessage.class);
        registerMessageType(HeartbeatResponse.class);
        registerMessageType(RegisterClientMessage.class);
        registerMessageType(RegisterWorkerMessage.class);
        registerMessageType(ResponseMessage.class);
        registerMessageType(RouteQueryMessage.class);
        registerMessageType(RouteResponse.class);
        registerMessageType(RpcRequest.class);
        registerMessageType(RpcResponse.class);
    }

    private static void registerMessageType(Class<? extends IMessage> messageClass) {
        messageLib.putIfAbsent(messageClass.getName(), messageClass);
    }

    /**
     * 序列化消息对象为字节数组
     * @param message 消息对象
     * @return 序列化后的字节数组
     */
    public static byte[] serialize(IMessage message) {
        try {
            // 使用UTF-8编码，确保中文不乱码
            // 添加@type字段以便反序列化时识别具体类型
            return JSON.toJSONString(message, JSONWriter.Feature.WriteClassName)
                    .getBytes(StandardCharsets.UTF_8);
        } catch (Exception e) {
            logger.error("Message serialization failed: {}", e.getMessage());
            throw new MessageException("消息序列化失败", e);
        }
    }

    /**
     * 反序列化字节数组为消息对象
     * @param bytes 字节数组
     * @param <T> 消息类型
     * @return 反序列化后的消息对象
     */
    public static <T extends IMessage> T deserialize(byte[] bytes) {
        try {
            if (bytes == null || bytes.length == 0) {
                throw new MessageException("Empty message data");
            }

            // 调试日志：打印原始数据
            if (logger.isTraceEnabled()) {
                logger.trace("Raw message data ({} bytes): {}", bytes.length,
                        new String(bytes, 0, Math.min(bytes.length, 1024), StandardCharsets.UTF_8));
            }

            String jsonStr = new String(bytes, StandardCharsets.UTF_8);
            JSONObject jsonObject = JSON.parseObject(jsonStr);

            // 通用枚举处理
            processEnums(jsonObject);

            String typeName = jsonObject.getString("@type");
            if (typeName == null) {
                throw new MessageException("Missing @type field in message");
            }

            Class<? extends IMessage> messageClass = messageLib.get(typeName);
            if (messageClass == null) {
                throw new MessageException("Unsupported message type: " + typeName);
            }

            return JSON.parseObject(jsonStr, (Type) messageClass);
        } catch (Exception e) {
            logger.error("Message deserialization failed. Data length: {}", bytes != null ? bytes.length : "null");
            throw new MessageException("消息反序列化失败", e);
        }
    }


    /**
     * 安全反序列化（带数据完整性校验）
     */
    public static <T extends IMessage> T safeDeserialize(byte[] bytes) {
        try {
            // 基础检查
            if (bytes == null || bytes.length == 0) {
                throw new MessageException("Empty message data");
            }

            String jsonStr = new String(bytes, StandardCharsets.UTF_8);

            // 校验JSON完整性
            if (!isJsonComplete(jsonStr)) {
                throw new MessageException("Incomplete JSON data");
            }

            // 原有反序列化逻辑
            return deserialize(bytes);
        } catch (Exception e) {
            logger.error("Message deserialization failed", e);
            throw new MessageException("消息反序列化失败", e);
        }
    }

    /**
     * 检查JSON数据是否完整
     */
    private static boolean isJsonComplete(String jsonStr) {
        // 简单检查大括号匹配
        int openBraces = countChars(jsonStr, '{');
        int closeBraces = countChars(jsonStr, '}');
        return openBraces == closeBraces;
    }

    private static int countChars(String str, char ch) {
        return (int) str.chars().filter(c -> c == ch).count();
    }

    private static void processEnums(JSONObject jsonObject) {
        for (Map.Entry<String, Object> entry : jsonObject.entrySet()) {
            Object value = entry.getValue();
            if (value instanceof String) {
                try {
                    // 尝试获取字段的实际类型
                    String fieldName = entry.getKey();
                    Class<?> fieldType = getFieldType(jsonObject, fieldName);

                    if (fieldType != null && fieldType.isEnum()) {
                        @SuppressWarnings("unchecked")
                        Class<? extends Enum> enumType = (Class<? extends Enum>) fieldType;
                        Enum<?> enumValue = Enum.valueOf(enumType, (String)value);
                        jsonObject.put(fieldName, enumValue);
                    }
                } catch (Exception e) {
                    // 不是枚举类型或转换失败，忽略
                }
            }
        }
    }

    private static Class<?> getFieldType(JSONObject jsonObject, String fieldName) {
        try {
            // 通过@type字段获取目标类
            String className = jsonObject.getString("@type");
            if (className != null) {
                Class<?> clazz = Class.forName(className);
                return clazz.getDeclaredField(fieldName).getType();
            }
        } catch (Exception e) {
            // 忽略异常
        }
        return null;
    }

    /**
     * 反序列化字节数组为指定类型的消息对象
     * @param bytes 字节数组
     * @param clazz 目标类型
     * @param <T> 消息类型
     * @return 反序列化后的消息对象
     */
    public static <T extends IMessage> T deserialize(byte[] bytes, Class<T> clazz) {
        try {
            String jsonStr = new String(bytes, StandardCharsets.UTF_8);
            T result = JSON.parseObject(jsonStr, clazz);

            if (!clazz.isInstance(result)) {
                throw new MessageException("Type mismatch: expected " + clazz.getName() +
                        ", got " + (result != null ? result.getClass().getName() : "null"));
            }

            return result;
        } catch (Exception e) {
            logger.error("Message deserialization failed for type {}: {}", clazz.getSimpleName(), e.getMessage());
            throw new MessageException("消息反序列化失败", e);
        }
    }

    // 自定义异常类
    public static class MessageException extends RuntimeException {
        public MessageException(String message) {
            super(message);
        }

        public MessageException(String message, Throwable cause) {
            super(message, cause);
        }
    }
}